use windows_sys::Win32::System::Threading::{WaitForSingleObject, PTP_CALLBACK_INSTANCE, PTP_WORK, TP_CALLBACK_ENVIRON_V3};

use windows_sys::Win32::System::LibraryLoader::{GetModuleHandleA, GetProcAddress};

use windows_sys::core::s;

use std::arch::global_asm;
use std::ptr::{null_mut};
use std::ffi::{c_void};


#[allow(non_camel_case_types)]
type PTP_WORK_CALLBACK = unsafe extern "system" fn(
    instance: PTP_CALLBACK_INSTANCE,
    context: *mut c_void, 
    work: PTP_WORK,
); 


// store loadliba addr globally
static mut P_LOAD_LIBRARY_A: usize = 0; 

#[unsafe (no_mangle)]
pub extern "C" fn get_load_library_a() -> usize {
    unsafe { P_LOAD_LIBRARY_A }
}

// workcallback
global_asm!(
    r#"
    .section .text
    .global WorkCallback
    WorkCallback:
        mov rcx, rdx
        xor rdx, rdx
        call get_load_library_a
        jmp rax
    "#
);

// declare WorkCallback as an external function
unsafe extern "C" {
    fn WorkCallback(
        instance: PTP_CALLBACK_INSTANCE,
        context: *mut std::ffi::c_void,
        work: PTP_WORK,
    );
}


fn main() {

    unsafe {
    
        let kernel32_handle = GetModuleHandleA(s!("kernel32"));
        
        let ntdll_handle = GetModuleHandleA(s!("ntdll"));

        // func add 

        P_LOAD_LIBRARY_A = GetProcAddress(kernel32_handle, s!("LoadLibraryA")).unwrap() as usize;

        let p_tp_alloc_work = GetProcAddress(ntdll_handle, s!("TpAllocWork")).unwrap();
        let p_tp_post_work = GetProcAddress(ntdll_handle, s!("TpPostWork")).unwrap();
        let p_tp_release_work = GetProcAddress(ntdll_handle, s!("TpReleaseWork")).unwrap();


        // loading wininet.dll using thread pool
        let lib_name = s!("wininet.dll");

        let mut work_return: PTP_WORK = 0;

        // call TpAllocWork
        let tp_alloc_work_fn: unsafe extern "system" fn(
            *mut PTP_WORK,
            PTP_WORK_CALLBACK,
            *mut c_void,
            *mut TP_CALLBACK_ENVIRON_V3,
        ) -> i32 = std::mem::transmute(p_tp_alloc_work);


        tp_alloc_work_fn(
            &mut work_return,

            std::mem::transmute::<
                unsafe extern "C" fn(PTP_CALLBACK_INSTANCE, *mut c_void, PTP_WORK),
                unsafe extern "system" fn(PTP_CALLBACK_INSTANCE, *mut c_void, PTP_WORK),
            >(WorkCallback),

            lib_name as *mut _,
            
            null_mut(),
        );

        // call TpPostWork
        let tp_post_work_fn: unsafe extern "system" fn(PTP_WORK) = std::mem::transmute(p_tp_post_work);
        tp_post_work_fn(work_return);


        // Call TpReleaseWork
        let tp_release_work_fn: unsafe extern "system" fn(PTP_WORK) = std::mem::transmute(p_tp_release_work);
        tp_release_work_fn(work_return);

        // Wait for the thread pool to process
        WaitForSingleObject(-1isize as _, 0x1000);
        
        // verify if our dll was loaded !

        let wininet_handle = GetModuleHandleA(lib_name);

        println!("hWininet: {:p}", wininet_handle as *const ());

        println!("[+] Press <Enter> to Exit");

        let mut sss = String::new();
        std::io::stdin().read_line(&mut sss).unwrap();

    }   

}
